#    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.


from copy import deepcopy
from nnunet.utilities.nd_softmax import softmax_helper
from torch import nn
import torch
import numpy as np
from nnunet.network_architecture.initialization import InitWeights_He
from nnunet.network_architecture.neural_network import SegmentationNetwork
import torch.nn.functional
from nnunet.network_architecture.custom_modules.yusongli.DFCAN import RCAB
from axial_attention import AxialAttention, AxialPositionalEmbedding
import copy
from vit_pytorch.vit_3d import ViT
from rich.pretty import pprint as print
import sys


# ! <<< open debug yusongli
class Change2(nn.Module):
    def __init__(self, conv_kwargs, transp_kwargs):
        super().__init__()
        conv_in_channels = conv_kwargs.get('in_channels')
        transp_in_channels = transp_kwargs.get('in_channels')
        conv_out_channels = conv_kwargs.get('out_channels')
        transp_out_channels = transp_kwargs.get('out_channels')
        self.conv = nn.Sequential(
            nn.Conv3d(**conv_kwargs),
            nn.BatchNorm3d(conv_out_channels, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True),
            nn.LeakyReLU(negative_slope=1e-2, inplace=True),
        )
        self.transp = nn.Sequential(
            nn.ConvTranspose3d(**transp_kwargs),
            nn.BatchNorm3d(transp_out_channels, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True),
            nn.LeakyReLU(negative_slope=1e-2, inplace=True),
        )
        self.out_up = nn.Conv3d(conv_in_channels * 2, conv_in_channels, kernel_size=1, stride=1)
        self.out_down = nn.Conv3d(transp_in_channels * 2, transp_in_channels, kernel_size=1, stride=1)

    def forward(self, x_up, x_down):
        # (1, 256, 12, 8, 10)
        x_up1 = self.conv(x_up)
        # (1, 128, 24, 16, 20)
        x_down1 = self.transp(x_down)

        x_up = torch.concat([x_up, x_down1], dim=1)
        x_down = torch.concat([x_down, x_up1], dim=1)

        x_up = self.out_up(x_up)
        x_down = self.out_down(x_down)

        return x_up, x_down


class AxialAttention3D(nn.Module):
    def __init__(
        self,
        n_features,
        emb_shape,
        heads,
        dim_heads,
        non_lin,
        non_lin_kwargs,
        bn=True,
        sum_axial_out=True,
        residual_attention=False,
    ):
        super().__init__()
        self.axial_attention = [
            AxialPositionalEmbedding(dim=n_features, shape=emb_shape),
            AxialAttention(
                dim=n_features,
                heads=heads,
                dim_heads=dim_heads,
                dim_index=1,
                num_dimensions=3,
                sum_axial_out=sum_axial_out,
            ),
        ]
        if bn:
            self.axial_attention += [nn.BatchNorm3d(n_features)]

        if non_lin is not None:
            self.axial_attention += [non_lin(**non_lin_kwargs)]
        self.axial_attention = nn.Sequential(*self.axial_attention)
        self.residual_attention = residual_attention

    def forward(self, x):
        if self.residual_attention:
            return x + self.axial_attention(x)
        else:
            return self.axial_attention(x)


class ArchBlock(nn.Module):
    def __init__(self, **kwargs):
        super().__init__()
        # RCAB
        dim = kwargs.pop('dim')
        input_channels = kwargs.pop('input_channels')
        hidden_channels = kwargs.pop('hidden_channels')

        # AxialAttention3D
        n_features = kwargs.pop('n_features')
        emb_shape = kwargs.pop('emb_shape')
        heads = kwargs.pop('heads')
        dim_heads = kwargs.pop('dim_heads')
        nonlin = kwargs.pop('nonlin')
        nonlin_kwargs = kwargs.pop('nonlin_kwargs')
        bn = kwargs.pop('bn')
        sum_axial_out = kwargs.pop('sum_axial_out')
        residual_attention = kwargs.pop('residual_attention')

        # self.rcab = RCAB(dim=dim, input_channels=input_channels, hidden_channels=hidden_channels)
        self.att = AxialAttention3D(n_features, emb_shape, heads, dim_heads, nonlin, nonlin_kwargs, bn=bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,)

    def forward(self, x):
        x0 = x
        # x = self.rcab(x)
        x = self.att(x)
        # x = x + self.rcab(x)
        # x = x + self.att(x)
        return x0 + x


class ArchBlock02(nn.Module):
    def __init__(self, **kwargs):
        super().__init__()

    def forward(self, x):
        return x


class ArchBlock0301(nn.Module):
    def __init__(self, **kwargs):
        super().__init__()
        # AxialAttention3D
        n_features = kwargs.pop('n_features')
        emb_shape = kwargs.pop('emb_shape')
        heads = kwargs.pop('heads')
        dim_heads = kwargs.pop('dim_heads')
        nonlin = kwargs.pop('nonlin')
        nonlin_kwargs = kwargs.pop('nonlin_kwargs')
        bn = kwargs.pop('bn')
        sum_axial_out = kwargs.pop('sum_axial_out')
        residual_attention = kwargs.pop('residual_attention')

        self.att = AxialAttention3D(
            n_features,
            emb_shape,
            heads,
            dim_heads,
            nonlin,
            nonlin_kwargs,
            bn=bn,
            sum_axial_out=sum_axial_out,
            residual_attention=residual_attention,
        )

    def forward(self, x):
        return self.att(x)


class ArchBlock08(nn.Module):
    def __init__(self, **kwargs):
        super().__init__()
        # AxialAttention3D
        n_features = kwargs.pop('n_features')
        emb_shape = kwargs.pop('emb_shape')
        heads = kwargs.pop('heads')
        dim_heads = kwargs.pop('dim_heads')
        nonlin = kwargs.pop('nonlin')
        nonlin_kwargs = kwargs.pop('nonlin_kwargs')
        bn = kwargs.pop('bn')
        sum_axial_out = kwargs.pop('sum_axial_out')
        residual_attention = kwargs.pop('residual_attention')
        import ipdb; ipdb.set_trace()  # ! debug yusongli
        # self.vit = ViT(
        #     image_size =
        # )

    def forward(self, x):
        return x
# ! >>> clos debug


class ConvDropoutNormNonlin(nn.Module):
    """
    fixes a bug in ConvDropoutNormNonlin where lrelu was used regardless of nonlin. Bad.
    """

    def __init__(
        self,
        input_channels,
        output_channels,
        conv_op=nn.Conv2d,
        conv_kwargs=None,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
    ):
        super(ConvDropoutNormNonlin, self).__init__()
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}
        if conv_kwargs is None:
            conv_kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1, 'dilation': 1, 'bias': True}

        self.nonlin_kwargs = nonlin_kwargs
        self.nonlin = nonlin
        self.dropout_op = dropout_op
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.conv_kwargs = conv_kwargs
        self.conv_op = conv_op
        self.norm_op = norm_op

        self.conv = self.conv_op(input_channels, output_channels, **self.conv_kwargs)
        if self.dropout_op is not None and self.dropout_op_kwargs['p'] is not None and self.dropout_op_kwargs['p'] > 0:
            self.dropout = self.dropout_op(**self.dropout_op_kwargs)
        else:
            self.dropout = None
        self.instnorm = self.norm_op(output_channels, **self.norm_op_kwargs)
        self.lrelu = self.nonlin(**self.nonlin_kwargs)

    def forward(self, x):
        x = self.conv(x)
        if self.dropout is not None:
            x = self.dropout(x)
        return self.lrelu(self.instnorm(x))


class ConvDropoutNonlinNorm(ConvDropoutNormNonlin):
    def forward(self, x):
        x = self.conv(x)
        if self.dropout is not None:
            x = self.dropout(x)
        return self.instnorm(self.lrelu(x))


class StackedConvLayers(nn.Module):
    def __init__(
        self,
        input_feature_channels,
        output_feature_channels,
        num_convs,
        conv_op=nn.Conv2d,
        conv_kwargs=None,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        first_stride=None,
        basic_block=ConvDropoutNormNonlin,
    ):
        '''
        stacks ConvDropoutNormLReLU layers. initial_stride will only be applied to first layer in the stack. The other parameters affect all layers
        :param input_feature_channels:
        :param output_feature_channels:
        :param num_convs:
        :param dilation:
        :param kernel_size:
        :param padding:
        :param dropout:
        :param initial_stride:
        :param conv_op:
        :param norm_op:
        :param dropout_op:
        :param inplace:
        :param neg_slope:
        :param norm_affine:
        :param conv_bias:
        '''
        self.input_channels = input_feature_channels
        self.output_channels = output_feature_channels

        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}
        if conv_kwargs is None:
            conv_kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1, 'dilation': 1, 'bias': True}

        self.nonlin_kwargs = nonlin_kwargs
        self.nonlin = nonlin
        self.dropout_op = dropout_op
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.conv_kwargs = conv_kwargs
        self.conv_op = conv_op
        self.norm_op = norm_op

        if first_stride is not None:
            self.conv_kwargs_first_conv = deepcopy(conv_kwargs)
            self.conv_kwargs_first_conv['stride'] = first_stride
        else:
            self.conv_kwargs_first_conv = conv_kwargs

        super().__init__()
        self.blocks = nn.Sequential(
            *(
                [
                    basic_block(
                        input_feature_channels,
                        output_feature_channels,
                        self.conv_op,
                        self.conv_kwargs_first_conv,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                    )
                ]
                + [
                    basic_block(
                        output_feature_channels,
                        output_feature_channels,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                    )
                    for _ in range(num_convs - 1)
                ]
            )
        )

    def forward(self, x):
        return self.blocks(x)


def print_module_training_status(module):
    if (
        isinstance(module, nn.Conv2d)
        or isinstance(module, nn.Conv3d)
        or isinstance(module, nn.Dropout3d)
        or isinstance(module, nn.Dropout2d)
        or isinstance(module, nn.Dropout)
        or isinstance(module, nn.InstanceNorm3d)
        or isinstance(module, nn.InstanceNorm2d)
        or isinstance(module, nn.InstanceNorm1d)
        or isinstance(module, nn.BatchNorm2d)
        or isinstance(module, nn.BatchNorm3d)
        or isinstance(module, nn.BatchNorm1d)
    ):
        print(str(module), module.training)


class Upsample(nn.Module):
    def __init__(self, size=None, scale_factor=None, mode='nearest', align_corners=False):
        super(Upsample, self).__init__()
        self.align_corners = align_corners
        self.mode = mode
        self.scale_factor = scale_factor
        self.size = size

    def forward(self, x):
        return nn.functional.interpolate(
            x, size=self.size, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners
        )


class NYUNet(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock02(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                ) if d not in [0] else None
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        # (30, 1, 24, 64, 80)
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)
            # (30, 32, 24, 64, 80)
            # (30, 64, 24, 32, 40)
            # (30, 128, 24, 16, 20)
            # (30, 256, 12, 8, 10)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        # (30, 320, 6, 4, 5)
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        # if self.arch[0] is None:
        #     skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        #     skips[1], _        = self.changes[1][0](skips[1], skips[2])
        # else:
        #     skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        #     skips[1], _        = self.changes[1][0](skips[1], skips[2])
        #     skips[0], _        = self.changes[0][0](skips[0], skips[1])
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]) if self.arch[d] is not None else skips[d])

        # * Change (back)
        # if self.arch[0] is None:
        #     skibs[1], skibs[2] = self.changes[1][1](skibs[1], skibs[2])
        #     _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # else:
        #     skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        #     _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        #     _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        skibs[1], skibs[2] = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet01(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet0101(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                # ! <<< open debug yusongli
                # ArchBlock0301(
                #     dim=3,
                #     input_channels=input_channels,
                #     hidden_channels=hidden_channels,
                #     n_features=output_features,
                #     emb_shape=emb_shape,
                #     heads=(heads * 2**d),
                #     dim_heads=(dim_heads * 2**d),
                #     nonlin=nonlin,
                #     nonlin_kwargs=nonlin_kwargs,
                #     bn=axial_bn,
                #     sum_axial_out=sum_axial_out,
                #     residual_attention=residual_attention,
                # )
                # ! ===
                AxialAttention3D(
                    output_features,
                    emb_shape,
                    heads * 2**d,
                    dim_heads * 2**d,
                    nonlin,
                    nonlin_kwargs,
                    axial_bn,
                    sum_axial_out,
                    residual_attention
                )
                # ! >>> clos debug
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            print(x.shape)
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet0102(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                nn.Sequential(
                        StackedConvLayers(
                        input_features,
                        output_features,
                        num_conv_per_stage,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        first_stride,
                        basic_block=basic_block,
                    ),
                    RCAB(
                        3,
                        output_features,
                        [output_features // 2, output_features // 2],
                    )
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                # ! <<< open debug yusongli
                # ArchBlock0301(
                #     dim=3,
                #     input_channels=input_channels,
                #     hidden_channels=hidden_channels,
                #     n_features=output_features,
                #     emb_shape=emb_shape,
                #     heads=(heads * 2**d),
                #     dim_heads=(dim_heads * 2**d),
                #     nonlin=nonlin,
                #     nonlin_kwargs=nonlin_kwargs,
                #     bn=axial_bn,
                #     sum_axial_out=sum_axial_out,
                #     residual_attention=residual_attention,
                # )
                # ! ===
                AxialAttention3D(
                    output_features,
                    emb_shape,
                    heads * 2**d,
                    dim_heads * 2**d,
                    nonlin,
                    nonlin_kwargs,
                    axial_bn,
                    sum_axial_out,
                    residual_attention
                )
                # ! >>> clos debug
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            # ! <<< open debug yusongli
            # nfeatures_from_skip = self.conv_blocks_context[
            #     -(2 + u)
            # ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            # ! ===
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ][0].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            # ! >>> clos debug
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    RCAB(
                        3,
                        nfeatures_from_skip,
                        [nfeatures_from_skip // 2, nfeatures_from_skip // 2],
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet0103(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                nn.Sequential(
                        StackedConvLayers(
                        input_features,
                        output_features,
                        num_conv_per_stage,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        first_stride,
                        basic_block=basic_block,
                    ),
                    RCAB(
                        3,
                        output_features,
                        [output_features // 2, output_features // 2],
                    )
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                # ! <<< open debug yusongli
                # ArchBlock0301(
                #     dim=3,
                #     input_channels=input_channels,
                #     hidden_channels=hidden_channels,
                #     n_features=output_features,
                #     emb_shape=emb_shape,
                #     heads=(heads * 2**d),
                #     dim_heads=(dim_heads * 2**d),
                #     nonlin=nonlin,
                #     nonlin_kwargs=nonlin_kwargs,
                #     bn=axial_bn,
                #     sum_axial_out=sum_axial_out,
                #     residual_attention=residual_attention,
                # )
                # ! ===
                AxialAttention3D(
                    output_features,
                    emb_shape,
                    heads * 2**d,
                    dim_heads * 2**d,
                    nonlin,
                    nonlin_kwargs,
                    axial_bn,
                    sum_axial_out,
                    residual_attention
                )
                # ! >>> clos debug
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            # ! <<< open debug yusongli
            # nfeatures_from_skip = self.conv_blocks_context[
            #     -(2 + u)
            # ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            # ! ===
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ][0].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            # ! >>> clos debug
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    RCAB(
                        3,
                        nfeatures_from_skip,
                        [nfeatures_from_skip // 2, nfeatures_from_skip // 2],
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    # ! <<< open debug yusongli
                    # self.conv_blocks_localization[ds][-1].output_channels,
                    # ! ===
                    self.conv_blocks_localization[ds][-2].output_channels,
                    # ! >>> clos debug
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet02(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock02(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet03(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ? params: middle deep supervision
        self.mid_deep_supervision = []
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock02(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            self.mid_deep_supervision.append(
                nn.Conv3d(input_channels, 2, kernel_size=1, stride=1)
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        self.mid_deep_supervision = nn.ModuleList(self.mid_deep_supervision)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            # ! <<< open debug yusongli
            # return tuple(
            #     [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            # )
            # ! ===
            return (
                [self.final_nonlin(i(j)) for i, j in zip(self.mid_deep_supervision, skips)],
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
            # ! >>> clos debug
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet0301(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ? params: arch deep supervision
        self.deep_supervision_archpre = []
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock0301(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            self.deep_supervision_archpre.append(
                nn.Conv3d(input_channels, 2, kernel_size=1, stride=1)
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # ! <<< open debug yusongli
        self.deep_supervision_arch = copy.deepcopy(self.deep_supervision_archpre)
        # ! >>> clos debug

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            # ! <<< open debug yusongli
            # n_features_after_tu_and_concat = nfeatures_from_skip * 2
            # ! ===
            n_features_after_tu_and_concat = nfeatures_from_skip * 3
            # ! >>> clos debug

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        self.ddpm = nn.ModuleList([nn.Conv3d(6, 2, kernel_size=1, stride=1) for _ in range(4)])
        # ! >>> clos debug

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        self.deep_supervision_archpre = nn.ModuleList(self.deep_supervision_archpre)
        self.deep_supervision_arch = nn.ModuleList(self.deep_supervision_arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skips_bak = []
        skibs = []
        skibs_bak = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            # ! <<< open debug yusongli
            skips_bak.append(x)
            # ! >>> clos debug
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            temp = self.arch[d](skips[d])
            skibs.append(temp)
            # ! <<< open debug yusongli
            skibs_bak.append(temp)
            # ! >>> clos debug

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            # x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((skips_bak[-(u + 1)], skibs[-(u + 1)], x), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            # ! <<< open debug yusongli
            # seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))
            # ! ===
            seg_outputs.append(self.seg_outputs[u](x))
            # ! >>> clos debug

        # ! <<< open debug yusongli
        result_deep_supervision_archpre = [i(j) for i, j in zip(self.deep_supervision_archpre, skips)]
        result_deep_supervision_arch = [i(j) for i,j in zip(self.deep_supervision_arch, skibs_bak)]
        result_seg_outputs = [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
        result_ddpm = [
                self.ddpm[i](
                    torch.cat((
                        result_deep_supervision_archpre[i],
                        result_deep_supervision_arch[i],
                        result_seg_outputs[i]
                    ), dim=1)
                )
            for i in range(len(result_seg_outputs))
        ]
        # ! >>> clos debug

        if self._deep_supervision and self.do_ds:
            # ! <<< open debug yusongli
            # return tuple(
            #     [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            # )
            # ! ===
            return (
                [self.final_nonlin(item) for item in result_deep_supervision_archpre],
                [self.final_nonlin(item) for item in result_deep_supervision_arch],
                [self.final_nonlin(item) for item in result_seg_outputs],  # ! There might be a bug if `self.upscale_logits_ops` is not `lambda x: x`, so this line needs to be checked in the future.
                [self.final_nonlin(item) for item in result_ddpm]
            )
            # ! >>> clos debug
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet0302(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ? params: arch deep supervision
        self.deep_supervision_archpre = []
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock0301(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            self.deep_supervision_archpre.append(
                nn.Conv3d(input_channels, 2, kernel_size=1, stride=1)
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # ! <<< open debug yusongli
        self.deep_supervision_arch = copy.deepcopy(self.deep_supervision_archpre)
        # ! >>> clos debug

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            # ! <<< open debug yusongli
            # n_features_after_tu_and_concat = nfeatures_from_skip * 2
            # ! ===
            # n_features_after_tu_and_concat = nfeatures_from_skip * 3
            # ! ===
            n_features_after_tu_and_concat = nfeatures_from_skip * 3 + 4
            # ! >>> clos debug

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # self.ddpm = nn.ModuleList([nn.Conv3d(6, 2, kernel_size=1, stride=1) for _ in range(4)])
        # ! >>> clos debug

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        self.deep_supervision_archpre = nn.ModuleList(self.deep_supervision_archpre)
        self.deep_supervision_arch = nn.ModuleList(self.deep_supervision_arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skips_bak = []
        skibs = []
        skibs_bak = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            # ! <<< open debug yusongli
            skips_bak.append(x)
            # ! >>> clos debug
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            temp = self.arch[d](skips[d])
            skibs.append(temp)
            # ! <<< open debug yusongli
            skibs_bak.append(temp)
            # ! >>> clos debug

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # ! <<< open debug yusongli
        result_deep_supervision_archpre = [i(j) for i, j in zip(self.deep_supervision_archpre, skips)]
        result_deep_supervision_arch = [i(j) for i,j in zip(self.deep_supervision_arch, skibs_bak)]
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            # x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! ===
            # x = torch.cat((skips_bak[-(u + 1)],skibs[-(u + 1)], x), dim=1)
            # ! ===
            x = torch.cat((
                skips_bak[-(u + 1)],
                result_deep_supervision_archpre[-(u + 1)],
                skibs[-(u + 1)],
                result_deep_supervision_arch[-(u + 1)],
                x
            ), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            # ! <<< open debug yusongli
            # seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))
            # ! ===
            seg_outputs.append(self.seg_outputs[u](x))
            # ! >>> clos debug

        # ! <<< open debug yusongli
        result_seg_outputs = [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
        # result_ddpm = [
        #         self.ddpm[i](
        #             torch.cat((
        #                 result_deep_supervision_archpre[i],
        #                 result_deep_supervision_arch[i],
        #                 result_seg_outputs[i]
        #             ), dim=1)
        #         )
        #     for i in range(len(result_seg_outputs))
        # ]
        # ! >>> clos debug

        if self._deep_supervision and self.do_ds:
            # ! <<< open debug yusongli
            # return tuple(
            #     [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            # )
            # ! ===
            return (
                [self.final_nonlin(item) for item in result_deep_supervision_archpre],
                [self.final_nonlin(item) for item in result_deep_supervision_arch],
                [self.final_nonlin(item) for item in result_seg_outputs],  # ! There might be a bug if `self.upscale_logits_ops` is not `lambda x: x`, so this line needs to be checked in the future.
                # [self.final_nonlin(item) for item in result_ddpm]
            )
            # ! >>> clos debug
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet0303(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ? params: arch deep supervision
        self.deep_supervision_archpre = []
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock0301(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            self.deep_supervision_archpre.append(
                nn.Conv3d(input_channels, 2, kernel_size=1, stride=1)
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # ! <<< open debug yusongli
        self.deep_supervision_arch = copy.deepcopy(self.deep_supervision_archpre)
        # ! >>> clos debug

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            # ! <<< open debug yusongli
            # n_features_after_tu_and_concat = nfeatures_from_skip * 2
            # ! ===
            n_features_after_tu_and_concat = nfeatures_from_skip * 3
            # ! >>> clos debug

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # self.ddpm = nn.ModuleList([nn.Conv3d(6, 2, kernel_size=1, stride=1) for _ in range(4)])
        # ! >>> clos debug

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        self.deep_supervision_archpre = nn.ModuleList(self.deep_supervision_archpre)
        self.deep_supervision_arch = nn.ModuleList(self.deep_supervision_arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skips_bak = []
        skibs = []
        skibs_bak = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            # ! <<< open debug yusongli
            skips_bak.append(x)
            # ! >>> clos debug
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            temp = self.arch[d](skips[d])
            skibs.append(temp)
            # ! <<< open debug yusongli
            skibs_bak.append(temp)
            # ! >>> clos debug

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # ! <<< open debug yusongli
        result_deep_supervision_archpre = [i(j) for i, j in zip(self.deep_supervision_archpre, skips)]
        result_deep_supervision_arch = [i(j) for i,j in zip(self.deep_supervision_arch, skibs_bak)]
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            # x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! ===
            # x = torch.cat((skips_bak[-(u + 1)],skibs[-(u + 1)], x), dim=1)
            # ! ===
            # ? Try to generate mask, and use the mask to teach network to focus the important parts.
            x = torch.cat((
                skips_bak[-(u + 1)] * result_deep_supervision_archpre[-(u + 1)].max(dim=1).indices.unsqueeze(dim=1) + skips_bak[-(u + 1)],
                skibs[-(u + 1)] * result_deep_supervision_arch[-(u + 1)].max(dim=1).indices.unsqueeze(dim=1) + skibs[-(u + 1)],
                x
            ), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            # ! <<< open debug yusongli
            # seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))
            # ! ===
            seg_outputs.append(self.seg_outputs[u](x))
            # ! >>> clos debug

        # ! <<< open debug yusongli
        result_seg_outputs = [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
        # result_ddpm = [
        #         self.ddpm[i](
        #             torch.cat((
        #                 result_deep_supervision_archpre[i],
        #                 result_deep_supervision_arch[i],
        #                 result_seg_outputs[i]
        #             ), dim=1)
        #         )
        #     for i in range(len(result_seg_outputs))
        # ]
        # ! >>> clos debug

        if self._deep_supervision and self.do_ds:
            # ! <<< open debug yusongli
            # return tuple(
            #     [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            # )
            # ! ===
            return (
                [self.final_nonlin(item) for item in result_deep_supervision_archpre],
                [self.final_nonlin(item) for item in result_deep_supervision_arch],
                [self.final_nonlin(item) for item in result_seg_outputs],  # ! There might be a bug if `self.upscale_logits_ops` is not `lambda x: x`, so this line needs to be checked in the future.
                # [self.final_nonlin(item) for item in result_ddpm]
            )
            # ! >>> clos debug
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet030301(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ? params: arch deep supervision
        self.deep_supervision_archpre = []
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock0301(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            self.deep_supervision_archpre.append(
                nn.Conv3d(input_channels, 2, kernel_size=1, stride=1)
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # ! <<< open debug yusongli
        self.deep_supervision_arch = copy.deepcopy(self.deep_supervision_archpre)
        # ! >>> clos debug

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            # ! <<< open debug yusongli
            # n_features_after_tu_and_concat = nfeatures_from_skip * 2
            # ! ===
            n_features_after_tu_and_concat = nfeatures_from_skip * 3
            # ! >>> clos debug

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # self.ddpm = nn.ModuleList([nn.Conv3d(6, 2, kernel_size=1, stride=1) for _ in range(4)])
        # ! >>> clos debug

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        self.deep_supervision_archpre = nn.ModuleList(self.deep_supervision_archpre)
        self.deep_supervision_arch = nn.ModuleList(self.deep_supervision_arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skips_bak = []
        skibs = []
        skibs_bak = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            # ! <<< open debug yusongli
            skips_bak.append(x)
            # ! >>> clos debug
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            temp = self.arch[d](skips[d])
            skibs.append(temp)
            # ! <<< open debug yusongli
            skibs_bak.append(temp)
            # ! >>> clos debug

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # ! <<< open debug yusongli
        result_deep_supervision_archpre = [i(j) for i, j in zip(self.deep_supervision_archpre, skips)]
        result_deep_supervision_arch = [i(j) for i,j in zip(self.deep_supervision_arch, skibs_bak)]
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            # x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! ===
            # x = torch.cat((skips_bak[-(u + 1)],skibs[-(u + 1)], x), dim=1)
            # ! ===
            # ? Try to generate mask, and use the mask to teach network to focus the important parts.
            x = torch.cat((
                skips_bak[-(u + 1)] * torch.nn.functional.softmax(result_deep_supervision_archpre[-(u + 1)], dim=1).max(dim=1).indices.unsqueeze(dim=1) + skips_bak[-(u + 1)],
                skibs[-(u + 1)] * torch.nn.functional.softmax(result_deep_supervision_arch[-(u + 1)], dim=1).max(dim=1).indices.unsqueeze(dim=1) + skibs[-(u + 1)],
                x
            ), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            # ! <<< open debug yusongli
            # seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))
            # ! ===
            seg_outputs.append(self.seg_outputs[u](x))
            # ! >>> clos debug

        # ! <<< open debug yusongli
        result_seg_outputs = [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
        # result_ddpm = [
        #         self.ddpm[i](
        #             torch.cat((
        #                 result_deep_supervision_archpre[i],
        #                 result_deep_supervision_arch[i],
        #                 result_seg_outputs[i]
        #             ), dim=1)
        #         )
        #     for i in range(len(result_seg_outputs))
        # ]
        # ! >>> clos debug

        if self._deep_supervision and self.do_ds:
            # ! <<< open debug yusongli
            # return tuple(
            #     [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            # )
            # ! ===
            return (
                [self.final_nonlin(item) for item in result_deep_supervision_archpre],
                [self.final_nonlin(item) for item in result_deep_supervision_arch],
                [self.final_nonlin(item) for item in result_seg_outputs],  # ! There might be a bug if `self.upscale_logits_ops` is not `lambda x: x`, so this line needs to be checked in the future.
                # [self.final_nonlin(item) for item in result_ddpm]
            )
            # ! >>> clos debug
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet0304(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ? params: arch deep supervision
        self.deep_supervision_archpre = []
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock0301(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            self.deep_supervision_archpre.append(
                nn.Conv3d(input_channels, 2, kernel_size=1, stride=1)
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # ! <<< open debug yusongli
        self.deep_supervision_arch = copy.deepcopy(self.deep_supervision_archpre)
        # ! >>> clos debug

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            # ! <<< open debug yusongli
            n_features_after_tu_and_concat = nfeatures_from_skip * 2
            # ! >>> clos debug

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # self.ddpm = nn.ModuleList([nn.Conv3d(6, 2, kernel_size=1, stride=1) for _ in range(4)])
        # ! >>> clos debug

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        self.deep_supervision_archpre = nn.ModuleList(self.deep_supervision_archpre)
        self.deep_supervision_arch = nn.ModuleList(self.deep_supervision_arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skips_bak = []
        skibs = []
        skibs_bak = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            # ! <<< open debug yusongli
            skips_bak.append(x)
            # ! >>> clos debug
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            temp = self.arch[d](skips[d])
            skibs.append(temp)
            # ! <<< open debug yusongli
            skibs_bak.append(temp)
            # ! >>> clos debug

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # ! <<< open debug yusongli
        result_deep_supervision_archpre = [i(j) for i, j in zip(self.deep_supervision_archpre, skips)]
        result_deep_supervision_arch = [i(j) for i,j in zip(self.deep_supervision_arch, skibs_bak)]
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            # x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! ===
            # x = torch.cat((skips_bak[-(u + 1)],skibs[-(u + 1)], x), dim=1)
            # ! ===
            # ? Try to generate mask, and use the mask to teach network to focus the important parts.
            # x = torch.cat((
            #         skips_bak[-(u + 1)] * result_deep_supervision_archpre[-(u + 1)].max(dim=1).indices.unsqueeze(dim=1) + skips_bak[-(u + 1)],
            #         skibs[-(u + 1)] * result_deep_supervision_arch[-(u + 1)].max(dim=1).indices.unsqueeze(dim=1) + skibs[-(u + 1)],
            #         x
            #     ), dim=1)
            # ! ===
            temp1 = torch.nn.functional.softmax(result_deep_supervision_archpre[-(u + 1)], dim=1).max(dim=1).indices.unsqueeze(dim=1)
            temp2 = torch.nn.functional.softmax(result_deep_supervision_arch[-(u + 1)], dim=1).max(dim=1).indices.unsqueeze(dim=1)
            temp3 = (temp1 == temp2) * temp1
            x = torch.cat((
                skibs[-(u + 1)] * temp3 + skibs[-(u + 1)],
                x
            ), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            # ! <<< open debug yusongli
            # seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))
            # ! ===
            seg_outputs.append(self.seg_outputs[u](x))
            # ! >>> clos debug

        # ! <<< open debug yusongli
        result_seg_outputs = [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
        # result_ddpm = [
        #         self.ddpm[i](
        #             torch.cat((
        #                 result_deep_supervision_archpre[i],
        #                 result_deep_supervision_arch[i],
        #                 result_seg_outputs[i]
        #             ), dim=1)
        #         )
        #     for i in range(len(result_seg_outputs))
        # ]
        # ! >>> clos debug

        if self._deep_supervision and self.do_ds:
            # ! <<< open debug yusongli
            # return tuple(
            #     [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            # )
            # ! ===
            return (
                [self.final_nonlin(item) for item in result_deep_supervision_archpre],
                [self.final_nonlin(item) for item in result_deep_supervision_arch],
                [self.final_nonlin(item) for item in result_seg_outputs],  # ! There might be a bug if `self.upscale_logits_ops` is not `lambda x: x`, so this line needs to be checked in the future.
                # [self.final_nonlin(item) for item in result_ddpm]
            )
            # ! >>> clos debug
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet04(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock02(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skips_bak = []
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            # ! <<< open debug yusongli
            skips_bak.append(x)
            # ! >>> clos debug
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            x = x * skips_bak[-(u + 1)]
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet05(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ? params: middle deep supervision
        self.mid_deep_supervision = []
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock02(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            self.mid_deep_supervision.append(
                nn.Conv3d(input_channels, 2, kernel_size=1, stride=1)
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        self.mid_deep_supervision = nn.ModuleList(self.mid_deep_supervision)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skips_bak = []
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            # ! <<< open debug yusongli
            skips_bak.append(x)
            # ! >>> clos debug
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            x = x * skips_bak[-(u + 1)]
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            # ! <<< open debug yusongli
            # return tuple(
            #     [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            # )
            # ! ===
            return (
                [self.final_nonlin(i(j)) for i, j in zip(self.mid_deep_supervision, skips)],
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
            # ! >>> clos debug
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet06(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        # ! <<< open debug yusongli
        self.conv_blocks_localization01 = []
        # ! >>> clos debug
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock02(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        # ! <<< open debug yusongli
        self.conv_blocks_localization01 = copy.deepcopy(self.conv_blocks_localization)
        # ! >>> clos debug

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        # ! <<< open debug yusongli
        self.conv_blocks_localization01 = nn.ModuleList(self.conv_blocks_localization01)
        # ! >>> clos debug
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skips_bak = []
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            # ! <<< open debug yusongli
            skips_bak.append(x)
            # ! >>> clos debug
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skips_bak[-(u + 1)]), dim=1)
            x = self.conv_blocks_localization01[u](x)
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet07(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock02(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            # ! <<< open debug yusongli
            # n_features_after_tu_and_concat = nfeatures_from_skip * 2
            # ! ===
            n_features_after_tu_and_concat = nfeatures_from_skip * 3
            # ! >>> clos debug

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skips_bak = []
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            # ! <<< open debug yusongli
            skips_bak.append(x)
            # ! >>> clos debug
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((skips_bak[-(u + 1)], x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet08(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                ArchBlock08(
                    dim=3, input_channels=input_channels, hidden_channels=hidden_channels,
                    n_features=output_features, emb_shape=emb_shape, heads=(heads * 2**d), dim_heads=(dim_heads * 2**d), nonlin=nonlin, nonlin_kwargs=nonlin_kwargs, bn=axial_bn, sum_axial_out=sum_axial_out, residual_attention=residual_attention,
                )
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            # ! <<< open debug yusongli
            # n_features_after_tu_and_concat = nfeatures_from_skip * 2
            # ! ===
            n_features_after_tu_and_concat = nfeatures_from_skip * 3
            # ! >>> clos debug

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skips_bak = []
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            # ! <<< open debug yusongli
            skips_bak.append(x)
            # ! >>> clos debug
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((skips_bak[-(u + 1)], x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp

class NYUNet09(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                # ! <<< open debug yusongli
                # ArchBlock0301(
                #     dim=3,
                #     input_channels=input_channels,
                #     hidden_channels=hidden_channels,
                #     n_features=output_features,
                #     emb_shape=emb_shape,
                #     heads=(heads * 2**d),
                #     dim_heads=(dim_heads * 2**d),
                #     nonlin=nonlin,
                #     nonlin_kwargs=nonlin_kwargs,
                #     bn=axial_bn,
                #     sum_axial_out=sum_axial_out,
                #     residual_attention=residual_attention,
                # )
                # ! ===
                AxialAttention3D(
                    output_features,
                    emb_shape,
                    heads * 2**d,
                    dim_heads * 2**d,
                    nonlin,
                    nonlin_kwargs,
                    axial_bn,
                    sum_axial_out,
                    residual_attention
                )
                # ! >>> clos debug
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        # skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        # skips[1], _        = self.changes[1][0](skips[1], skips[2])
        # skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        # skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        # _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        # _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            # x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet1001(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                # ! <<< open debug yusongli
                # ArchBlock0301(
                #     dim=3,
                #     input_channels=input_channels,
                #     hidden_channels=hidden_channels,
                #     n_features=output_features,
                #     emb_shape=emb_shape,
                #     heads=(heads * 2**d),
                #     dim_heads=(dim_heads * 2**d),
                #     nonlin=nonlin,
                #     nonlin_kwargs=nonlin_kwargs,
                #     bn=axial_bn,
                #     sum_axial_out=sum_axial_out,
                #     residual_attention=residual_attention,
                # )
                # ! ===
                AxialAttention3D(
                    output_features,
                    emb_shape,
                    heads * 2**d,
                    dim_heads * 2**d,
                    nonlin,
                    nonlin_kwargs,
                    axial_bn,
                    sum_axial_out,
                    residual_attention
                )
                # ! >>> clos debug
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[0], skips[1] = self.changes[0][0](skips[0], skips[1])
        _, skips[2] = self.changes[1][0](skips[1], skips[2])
        _, skips[3] = self.changes[2][0](skips[2], skips[3])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        # skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        # _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        # _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet1002(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                # ! <<< open debug yusongli
                # ArchBlock0301(
                #     dim=3,
                #     input_channels=input_channels,
                #     hidden_channels=hidden_channels,
                #     n_features=output_features,
                #     emb_shape=emb_shape,
                #     heads=(heads * 2**d),
                #     dim_heads=(dim_heads * 2**d),
                #     nonlin=nonlin,
                #     nonlin_kwargs=nonlin_kwargs,
                #     bn=axial_bn,
                #     sum_axial_out=sum_axial_out,
                #     residual_attention=residual_attention,
                # )
                # ! ===
                AxialAttention3D(
                    output_features,
                    emb_shape,
                    heads * 2**d,
                    dim_heads * 2**d,
                    nonlin,
                    nonlin_kwargs,
                    axial_bn,
                    sum_axial_out,
                    residual_attention
                )
                # ! >>> clos debug
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> chttps://bava8u2znaj6bdzzjnfb.wgetcloud.online/link/yeXk009sAdypiiElos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        # skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        # _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        # _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet1003(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                # ! <<< open debug yusongli
                # ArchBlock0301(
                #     dim=3,
                #     input_channels=input_channels,
                #     hidden_channels=hidden_channels,
                #     n_features=output_features,
                #     emb_shape=emb_shape,
                #     heads=(heads * 2**d),
                #     dim_heads=(dim_heads * 2**d),
                #     nonlin=nonlin,
                #     nonlin_kwargs=nonlin_kwargs,
                #     bn=axial_bn,
                #     sum_axial_out=sum_axial_out,
                #     residual_attention=residual_attention,
                # )
                # ! ===
                AxialAttention3D(
                    output_features,
                    emb_shape,
                    heads * 2**d,
                    dim_heads * 2**d,
                    nonlin,
                    nonlin_kwargs,
                    axial_bn,
                    sum_axial_out,
                    residual_attention
                )
                # ! >>> clos debug
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        # skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        # skips[1], _        = self.changes[1][0](skips[1], skips[2])
        # skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet1004(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                # ! <<< open debug yusongli
                # ArchBlock0301(
                #     dim=3,
                #     input_channels=input_channels,
                #     hidden_channels=hidden_channels,
                #     n_features=output_features,
                #     emb_shape=emb_shape,
                #     heads=(heads * 2**d),
                #     dim_heads=(dim_heads * 2**d),
                #     nonlin=nonlin,
                #     nonlin_kwargs=nonlin_kwargs,
                #     bn=axial_bn,
                #     sum_axial_out=sum_axial_out,
                #     residual_attention=residual_attention,
                # )
                # ! ===
                AxialAttention3D(
                    output_features,
                    emb_shape,
                    heads * 2**d,
                    dim_heads * 2**d,
                    nonlin,
                    nonlin_kwargs,
                    axial_bn,
                    sum_axial_out,
                    residual_attention
                )
                # ! >>> clos debug
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        # skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        # skips[1], _        = self.changes[1][0](skips[1], skips[2])
        # skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[2], skibs[3] = self.changes[2][1](skibs[2], skibs[3])
        skibs[1], _ = self.changes[1][1](skibs[1], skibs[2])
        skibs[0], _ = self.changes[0][1](skibs[0], skibs[1])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet1005(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                # ! <<< open debug yusongli
                # ArchBlock0301(
                #     dim=3,
                #     input_channels=input_channels,
                #     hidden_channels=hidden_channels,
                #     n_features=output_features,
                #     emb_shape=emb_shape,
                #     heads=(heads * 2**d),
                #     dim_heads=(dim_heads * 2**d),
                #     nonlin=nonlin,
                #     nonlin_kwargs=nonlin_kwargs,
                #     bn=axial_bn,
                #     sum_axial_out=sum_axial_out,
                #     residual_attention=residual_attention,
                # )
                # ! ===
                AxialAttention3D(
                    output_features,
                    emb_shape,
                    heads * 2**d,
                    dim_heads * 2**d,
                    nonlin,
                    nonlin_kwargs,
                    axial_bn,
                    sum_axial_out,
                    residual_attention
                )
                # ! >>> clos debug
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[0], skips[1] = self.changes[0][0](skips[0], skips[1])
        _, skips[2] = self.changes[1][0](skips[1], skips[2])
        _, skips[3] = self.changes[2][0](skips[2], skips[3])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[2], skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        skibs[1], _ = self.changes[1][1](skibs[1], skibs[2])
        skibs[0], _ = self.changes[0][1](skibs[0], skibs[1])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet1006(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                # ! <<< open debug yusongli
                # ArchBlock0301(
                #     dim=3,
                #     input_channels=input_channels,
                #     hidden_channels=hidden_channels,
                #     n_features=output_features,
                #     emb_shape=emb_shape,
                #     heads=(heads * 2**d),
                #     dim_heads=(dim_heads * 2**d),
                #     nonlin=nonlin,
                #     nonlin_kwargs=nonlin_kwargs,
                #     bn=axial_bn,
                #     sum_axial_out=sum_axial_out,
                #     residual_attention=residual_attention,
                # )
                # ! ===
                AxialAttention3D(
                    output_features,
                    emb_shape,
                    heads * 2**d,
                    dim_heads * 2**d,
                    nonlin,
                    nonlin_kwargs,
                    axial_bn,
                    sum_axial_out,
                    residual_attention
                )
                # ! >>> clos debug
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[0], skips[1] = self.changes[0][0](skips[0], skips[1])
        _, skips[2] = self.changes[1][0](skips[1], skips[2])
        _, skips[3] = self.changes[2][0](skips[2], skips[3])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], skibs[1] = self.changes[0][1](skibs[0], skibs[1])
        _, skibs[2]        = self.changes[1][1](skibs[1], skibs[2])
        _, skibs[3]        = self.changes[2][1](skibs[2], skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet1007(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                # ! <<< open debug yusongli
                # ArchBlock0301(
                #     dim=3,
                #     input_channels=input_channels,
                #     hidden_channels=hidden_channels,
                #     n_features=output_features,
                #     emb_shape=emb_shape,
                #     heads=(heads * 2**d),
                #     dim_heads=(dim_heads * 2**d),
                #     nonlin=nonlin,
                #     nonlin_kwargs=nonlin_kwargs,
                #     bn=axial_bn,
                #     sum_axial_out=sum_axial_out,
                #     residual_attention=residual_attention,
                # )
                # ! ===
                AxialAttention3D(
                    output_features,
                    emb_shape,
                    heads * 2**d,
                    dim_heads * 2**d,
                    nonlin,
                    nonlin_kwargs,
                    axial_bn,
                    sum_axial_out,
                    residual_attention
                )
                # ! >>> clos debug
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        skips[2], skips[3] = self.changes[2][0](skips[2], skips[3])
        skips[1], _        = self.changes[1][0](skips[1], skips[2])
        skips[0], _        = self.changes[0][0](skips[0], skips[1])

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[2], skibs[3] = self.changes[2][1](skibs[2], skibs[3])
        skibs[1], _ = self.changes[1][1](skibs[1], skibs[2])
        skibs[0], _ = self.changes[0][1](skibs[0], skibs[1])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp


class NYUNet11(SegmentationNetwork):
    DEFAULT_BATCH_SIZE_3D = 2
    DEFAULT_PATCH_SIZE_3D = (64, 192, 160)
    SPACING_FACTOR_BETWEEN_STAGES = 2
    BASE_NUM_FEATURES_3D = 30
    MAX_NUMPOOL_3D = 999
    MAX_NUM_FILTERS_3D = 320

    DEFAULT_PATCH_SIZE_2D = (256, 256)
    BASE_NUM_FEATURES_2D = 30
    DEFAULT_BATCH_SIZE_2D = 50
    MAX_NUMPOOL_2D = 999
    MAX_FILTERS_2D = 480

    use_this_for_batch_size_computation_2D = 19739648
    use_this_for_batch_size_computation_3D = 520000000  # 505789440

    def __init__(
        self,
        input_channels,
        base_num_features,
        num_classes,
        num_pool,
        num_conv_per_stage=2,
        feat_map_mul_on_downscale=2,
        conv_op=nn.Conv2d,
        norm_op=nn.BatchNorm2d,
        norm_op_kwargs=None,
        dropout_op=nn.Dropout2d,
        dropout_op_kwargs=None,
        nonlin=nn.LeakyReLU,
        nonlin_kwargs=None,
        deep_supervision=True,
        dropout_in_localization=False,
        final_nonlin=softmax_helper,
        weightInitializer=InitWeights_He(1e-2),
        pool_op_kernel_sizes=None,
        conv_kernel_sizes=None,
        upscale_logits=False,
        convolutional_pooling=False,
        convolutional_upsampling=False,
        max_num_features=None,
        basic_block=ConvDropoutNormNonlin,
        seg_output_use_bias=False,
        # ! <<< open debug yusongli
        # ? Useful for Arch and Change
        heads=1,
        dim_heads=4,
        axial_bn=True,
        sum_axial_out=True,
        residual_attention=True,
        # ? Not Useful
        encoder_scale=1,
        axial_attention=False,
        volume_shape=(128, 128, 128),
        no_attention=[0],
        # ! >>> clos debug
    ):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> f.isensee@dkfz.de
        """
        super().__init__()
        self.convolutional_upsampling = convolutional_upsampling
        self.convolutional_pooling = convolutional_pooling
        self.upscale_logits = upscale_logits
        if nonlin_kwargs is None:
            nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        if dropout_op_kwargs is None:
            dropout_op_kwargs = {'p': 0.5, 'inplace': True}
        if norm_op_kwargs is None:
            norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1}

        self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True}

        self.nonlin = nonlin
        self.nonlin_kwargs = nonlin_kwargs
        self.dropout_op_kwargs = dropout_op_kwargs
        self.norm_op_kwargs = norm_op_kwargs
        self.weightInitializer = weightInitializer
        self.conv_op = conv_op
        self.norm_op = norm_op
        self.dropout_op = dropout_op
        self.num_classes = num_classes
        self.final_nonlin = final_nonlin
        self._deep_supervision = deep_supervision
        self.do_ds = deep_supervision

        if conv_op == nn.Conv2d:
            upsample_mode = 'bilinear'
            pool_op = nn.MaxPool2d
            transpconv = nn.ConvTranspose2d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3)] * (num_pool + 1)
        elif conv_op == nn.Conv3d:
            upsample_mode = 'trilinear'
            pool_op = nn.MaxPool3d
            transpconv = nn.ConvTranspose3d
            if pool_op_kernel_sizes is None:
                pool_op_kernel_sizes = [(2, 2, 2)] * num_pool
            if conv_kernel_sizes is None:
                conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1)
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(conv_op))

        self.input_shape_must_be_divisible_by = np.prod(pool_op_kernel_sizes, 0, dtype=np.int64)
        self.pool_op_kernel_sizes = pool_op_kernel_sizes
        self.conv_kernel_sizes = conv_kernel_sizes

        self.conv_pad_sizes = []
        for krnl in self.conv_kernel_sizes:
            self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl])

        if max_num_features is None:
            if self.conv_op == nn.Conv3d:
                self.max_num_features = self.MAX_NUM_FILTERS_3D
            else:
                self.max_num_features = self.MAX_FILTERS_2D
        else:
            self.max_num_features = max_num_features

        self.conv_blocks_context = []
        self.conv_blocks_localization = []
        self.td = []
        self.tu = []
        self.seg_outputs = []

        output_features = base_num_features
        input_features = input_channels

        # ! <<< open debug yusongli
        # ? Params: Change
        changes_kwargs = [
            [
                {'in_channels': 32, 'out_channels': 64, 'kernel_size': (1, 3, 3), 'stride': (1, 2, 2), 'padding': (0, 1, 1)},
                {'in_channels': 64, 'out_channels': 32, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 64, 'out_channels': 128, 'kernel_size': (3, 3, 3), 'stride': (1, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 128, 'out_channels': 64, 'kernel_size': (1, 2, 2), 'stride': (1, 2, 2)}
            ],
            [
                {'in_channels': 128, 'out_channels': 256, 'kernel_size': (3, 3, 3), 'stride': (2, 2, 2), 'padding': (1, 1, 1)},
                {'in_channels': 256, 'out_channels': 128, 'kernel_size': (2, 2, 2), 'stride': (2, 2, 2)}
            ]
        ]
        ChangeN = Change2
        self.changes = nn.ModuleList([
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[0][0],
                    transp_kwargs=changes_kwargs[0][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[1][0],
                    transp_kwargs=changes_kwargs[1][1],
                ),
            ]),
            nn.ModuleList([
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
                ChangeN(
                    conv_kwargs=changes_kwargs[2][0],
                    transp_kwargs=changes_kwargs[2][1],
                ),
            ])
        ])

        # ? Params: Arch
        self.arch = []
        ysl_emb_shape = [  # ? CZ2
            np.array([24, 80, 96]).astype(np.int16),
            np.array([24, 40, 48]).astype(np.int16),
            np.array([24, 20, 24]).astype(np.int16),
            np.array([12, 10, 12]).astype(np.int16),
            np.array([6, 5, 6]).astype(np.int16),
        ]
        ysl_input_channels = [32, 64, 128, 256]  # ? Z2
        # ! >>> clos debug

        for d in range(num_pool):
            # determine the first stride
            if d != 0 and self.convolutional_pooling:
                first_stride = pool_op_kernel_sizes[d - 1]
            else:
                first_stride = None

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[d]
            # add convolutions
            # * Down Sample Blocks
            self.conv_blocks_context.append(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                )
            )

            # ! <<< open debug yusongli
            # * Arch Blocks
            input_channels = ysl_input_channels[d]
            hidden_channels = [input_channels, input_channels]
            emb_shape = ysl_emb_shape[d]
            self.arch.append(
                # ! <<< open debug yusongli
                # ArchBlock0301(
                #     dim=3,
                #     input_channels=input_channels,
                #     hidden_channels=hidden_channels,
                #     n_features=output_features,
                #     emb_shape=emb_shape,
                #     heads=(heads * 2**d),
                #     dim_heads=(dim_heads * 2**d),
                #     nonlin=nonlin,
                #     nonlin_kwargs=nonlin_kwargs,
                #     bn=axial_bn,
                #     sum_axial_out=sum_axial_out,
                #     residual_attention=residual_attention,
                # )
                # ! ===
                AxialAttention3D(
                    output_features,
                    emb_shape,
                    heads * 2**d,
                    dim_heads * 2**d,
                    nonlin,
                    nonlin_kwargs,
                    axial_bn,
                    sum_axial_out,
                    residual_attention
                )
                # ! >>> clos debug
            )
            # ! >>> clos debug

            # * Max Pool Blocks
            if not self.convolutional_pooling:
                self.td.append(pool_op(pool_op_kernel_sizes[d]))
            input_features = output_features
            output_features = int(np.round(output_features * feat_map_mul_on_downscale))

            output_features = min(output_features, self.max_num_features)

        # now the bottleneck.
        # determine the first stride
        if self.convolutional_pooling:
            first_stride = pool_op_kernel_sizes[-1]
        else:
            first_stride = None

        # the output of the last conv must match the number of features from the skip connection if we are not using
        # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be
        # done by the transposed conv
        if self.convolutional_upsampling:
            final_num_features = output_features
        else:
            final_num_features = self.conv_blocks_context[-1].output_channels

        self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
        self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool]

        # * Bottleneck Blocks
        self.conv_blocks_context.append(
            nn.Sequential(
                StackedConvLayers(
                    input_features,
                    output_features,
                    num_conv_per_stage - 1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    first_stride,
                    basic_block=basic_block,
                ),
                StackedConvLayers(
                    output_features,
                    final_num_features,
                    1,
                    self.conv_op,
                    self.conv_kwargs,
                    self.norm_op,
                    self.norm_op_kwargs,
                    self.dropout_op,
                    self.dropout_op_kwargs,
                    self.nonlin,
                    self.nonlin_kwargs,
                    basic_block=basic_block,
                ),
            )
        )

        # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here
        if not dropout_in_localization:
            old_dropout_p = self.dropout_op_kwargs['p']
            self.dropout_op_kwargs['p'] = 0.0

        # now lets build the localization pathway
        for u in range(num_pool):
            nfeatures_from_down = final_num_features
            nfeatures_from_skip = self.conv_blocks_context[
                -(2 + u)
            ].output_channels  # self.conv_blocks_context[-1] is bottleneck, so start with -2
            n_features_after_tu_and_concat = nfeatures_from_skip * 2

            # the first conv reduces the number of features to match those of skip
            # the following convs work on that number of features
            # if not convolutional upsampling then the final conv reduces the num of features again
            if u != num_pool - 1 and not self.convolutional_upsampling:
                final_num_features = self.conv_blocks_context[-(3 + u)].output_channels
            else:
                final_num_features = nfeatures_from_skip

            # * Up Sample Blocks
            if not self.convolutional_upsampling:
                self.tu.append(Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode))
            else:
                self.tu.append(
                    transpconv(
                        nfeatures_from_down,
                        nfeatures_from_skip,
                        pool_op_kernel_sizes[-(u + 1)],
                        pool_op_kernel_sizes[-(u + 1)],
                        bias=False,
                    )
                )

            self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[-(u + 1)]
            self.conv_kwargs['padding'] = self.conv_pad_sizes[-(u + 1)]

            # * Merge Skip and Down Blocks
            self.conv_blocks_localization.append(
                nn.Sequential(
                    StackedConvLayers(
                        n_features_after_tu_and_concat,
                        nfeatures_from_skip,
                        num_conv_per_stage - 1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                    StackedConvLayers(
                        nfeatures_from_skip,
                        final_num_features,
                        1,
                        self.conv_op,
                        self.conv_kwargs,
                        self.norm_op,
                        self.norm_op_kwargs,
                        self.dropout_op,
                        self.dropout_op_kwargs,
                        self.nonlin,
                        self.nonlin_kwargs,
                        basic_block=basic_block,
                    ),
                )
            )

        for ds in range(len(self.conv_blocks_localization)):
            self.seg_outputs.append(
                conv_op(
                    self.conv_blocks_localization[ds][-1].output_channels,
                    num_classes,
                    1,
                    1,
                    0,
                    1,
                    1,
                    seg_output_use_bias,
                )
            )

        self.upscale_logits_ops = []
        cum_upsample = np.cumprod(np.vstack(pool_op_kernel_sizes), axis=0)[::-1]
        for usl in range(num_pool - 1):
            if self.upscale_logits:
                self.upscale_logits_ops.append(
                    Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), mode=upsample_mode)
                )
            else:
                self.upscale_logits_ops.append(lambda x: x)

        if not dropout_in_localization:
            self.dropout_op_kwargs['p'] = old_dropout_p

        # register all modules properly
        self.conv_blocks_localization = nn.ModuleList(self.conv_blocks_localization)
        self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context)
        self.td = nn.ModuleList(self.td)
        self.tu = nn.ModuleList(self.tu)
        self.seg_outputs = nn.ModuleList(self.seg_outputs)
        if self.upscale_logits:
            self.upscale_logits_ops = nn.ModuleList(
                self.upscale_logits_ops
            )  # lambda x:x is not a Module so we need to distinguish here

        # ! <<< open debug yusongli
        # * Device
        self.arch = nn.ModuleList(self.arch)
        # ! >>> clos debug

        if self.weightInitializer is not None:
            self.apply(self.weightInitializer)
            # self.apply(print_module_training_status)

    def forward(self, x):
        skips = []
        seg_outputs = []
        # ! <<< open debug yusongli
        skibs = []
        # ! >>> clos debug

        # * Down Sample
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)

            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        # * Bottleneck
        x = self.conv_blocks_context[-1](x)

        # ! <<< open debug yusongli
        # * Change (front)
        temp, skips[3] = self.changes[2][0](skips[2], skips[3])
        temp, skips[2] = self.changes[1][0](skips[1], temp)
        skips[0], skips[1] =self.changes[0][0](skips[0], temp)

        # * Arch
        for d in range(len(skips)):
            skibs.append(self.arch[d](skips[d]))

        # * Change (back)
        skibs[0], temp = self.changes[0][1](skibs[0], skibs[1])
        skibs[1], temp = self.changes[1][1](temp, skibs[2])
        skibs[2], skibs[3] = self.changes[2][1](temp, skibs[3])
        # ! >>> clos debug

        # * Up Sample
        for u in range(len(self.tu)):
            x = self.tu[u](x)
            # ! <<< open debug yusongli
            # x = torch.cat((x, skips[-(u + 1)]), dim=1)
            # ! ===
            x = torch.cat((x, skibs[-(u + 1)]), dim=1)
            # ! >>> clos debug
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))

        if self._deep_supervision and self.do_ds:
            return tuple(
                [seg_outputs[-1]] + [i(j) for i, j in zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]
            )
        else:
            return seg_outputs[-1]

    @staticmethod
    def compute_approx_vram_consumption(
        patch_size,
        num_pool_per_axis,
        base_num_features,
        max_num_features,
        num_modalities,
        num_classes,
        pool_op_kernel_sizes,
        deep_supervision=False,
        conv_per_stage=2,
    ):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param deep_supervision:
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :param num_modalities:
        :param num_classes:
        :param pool_op_kernel_sizes:
        :return:
        """
        if not isinstance(num_pool_per_axis, np.ndarray):
            num_pool_per_axis = np.array(num_pool_per_axis)

        npool = len(pool_op_kernel_sizes)

        map_size = np.array(patch_size)
        tmp = np.int64(
            (conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features
            + num_modalities * np.prod(map_size, dtype=np.int64)
            + num_classes * np.prod(map_size, dtype=np.int64)
        )

        num_feat = base_num_features

        for p in range(npool):
            for pi in range(len(num_pool_per_axis)):
                map_size[pi] /= pool_op_kernel_sizes[p][pi]
            num_feat = min(num_feat * 2, max_num_features)
            num_blocks = (
                (conv_per_stage * 2 + 1) if p < (npool - 1) else conv_per_stage
            )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
            tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
            if deep_supervision and p < (npool - 2):
                tmp += np.prod(map_size, dtype=np.int64) * num_classes
            # print(p, map_size, num_feat, tmp)
        return tmp
